# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
import datetime

from coin.base.datetime_util import to_datetime
from coin.base.hostname import get_hostname
from coin.base.timestamp import get_timestamp
from coin.exchange.base.proto_logger import (ProtoLogProvider, to_date_str, to_time_str)
from coin.proto.coin_executor_pb2 import (AggressiveExecutorConfig, PassiveExecutorConfig)
from coin.proto.coin_feed_pb2 import BestBookEntry
from coin.proto.coin_request_pb2 import (AccountRequestProto,
                                         StrategyRequestProto,
                                         ProcessInfoProto)
from coin.proto.coin_log_pb2 import (
    LogFileResponseProto,)
from coin.proto.coin_query_pb2 import (AccountBalance, AccountPosition, CurrencyBalance)
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.flow.archive import ensure_dir_exists
from coin.flow.topic_record import TopicRecordWriter

default_log_root = os.path.expanduser('~/data/strat_proto_log')


def gen_strategy_log_dir(*, log_root, trading_date, machine):
  date_str = trading_date.strftime('%Y%m%d')
  log_dir = os.path.join(log_root, machine, date_str)
  return log_dir


def compose_file_name(request: StrategyRequestProto, time_str: str, timestamp):
  assert type(request) == StrategyRequestProto
  log_name = '%s.%s.%s' % (request.strategy_name, time_str, timestamp)
  return log_name


# <strategy_name>.<YYmmDD-HHMMSS>.<timestamp_nano_sec>
# eg: test_mm_ma_strategy2.20180611-084408.1528706648659898000
def parse_file_name(file_name):
  gz_removed = file_name.replace('.gz', '')
  exploded = gz_removed.split('.')
  if len(exploded) == 3:
    timestamp = exploded[-1]
    time_str = exploded[-2]
    strategy_name = exploded[-3]
  elif len(exploded) == 4:
    # This is the case when synced log file is incomplete. There will be a
    # random suffix. Invalid and should be skipped.
    return None
  else:
    raise ValueError('Wrong file name: %s' % file_name)

  _date, _time = time_str.split('-')
  strat_req = StrategyRequestProto(strategy_name=strategy_name)
  return LogFileResponseProto(type=LogFileResponseProto.STRAT_LOG,
                              log_file_name=file_name,
                              date=_date,
                              time=_time,
                              timestamp=timestamp,
                              strat_req=strat_req)


def enumerate_logs(request: StrategyRequestProto, trading_date, log_dir):
  assert type(request) == StrategyRequestProto
  trading_date = trading_date.strftime('%Y%m%d')
  files_dict = {}
  if not os.path.isdir(log_dir):
    return []
  for file_name in os.listdir(log_dir):
    resp = parse_file_name(file_name)
    if resp is None:
      continue
    if (resp.strat_req.strategy_name == request.strategy_name and resp.date == trading_date):
      files_dict[resp.timestamp] = os.path.join(log_dir, file_name)
  return [files_dict[key] for key in sorted(files_dict.keys())]


class StrategyLogProvider(ProtoLogProvider):
  def __init__(self, request: StrategyRequestProto,
               log_root=None, machine=None):
    assert type(request) == StrategyRequestProto
    self._request = request
    self._log_root = log_root or default_log_root
    self._machine = machine or get_hostname()
    self._writer = None
    self._process = ProcessInfoProto(machine=self._machine, pid=os.getpid())
    super().__init__()

    self._balance_cache = None
    self._position_cache = None

  @property
  def request(self):
    return self._request

  def _gen_log_path_and_ensure(self, timestamp):
    log_root = self._log_root
    machine = self._machine
    timestamp = timestamp if timestamp is not None else get_timestamp()
    trading_date = to_datetime(timestamp)
    date_str = trading_date.strftime('%Y%m%d')
    time_str = to_time_str(timestamp)
    log_name = compose_file_name(request=self.request, time_str=time_str, timestamp=timestamp)

    log_dir = gen_strategy_log_dir(
        log_root=log_root, trading_date=trading_date, machine=machine)
    ensure_dir_exists(log_dir)
    log_path = os.path.join(log_dir, log_name)

    return LogFileResponseProto(type=LogFileResponseProto.STRAT_LOG,
                                log_file_name=log_path,
                                date=date_str,
                                time=time_str,
                                timestamp=str(timestamp),
                                strat_req=self.request)

  def gen_writer(self, timestamp):
    log_file_name = self._gen_log_path_and_ensure(timestamp).log_file_name
    if self._writer is None or self._writer.filepath != log_file_name:
      self._writer = TopicRecordWriter(log_file_name)
    return self._writer

  def write_position(self, protobuf, *,
                     flush: bool = True, timestamp=None):
    assert isinstance(protobuf, AccountPosition)
    log_proto = StrategyLog(type=StrategyLog.POSITION, position=protobuf)
    self.write(log_proto, flush, timestamp)

  def write_balance(self, protobuf, *,
                    flush: bool = True, timestamp=None):
    assert isinstance(protobuf, AccountBalance)
    log_proto = StrategyLog(type=StrategyLog.BALANCE, balance=protobuf)
    self.write(log_proto, flush, timestamp)

  def write_reserve(self, protobuf, *,
                    account_request=None, flush: bool = True, timestamp=None):
    assert isinstance(protobuf, AccountBalance)
    assert account_request is None or \
        isinstance(account_request, AccountRequestProto)
    log_proto = StrategyLog(type=StrategyLog.RESERVE,
                            reserve=protobuf,
                            account_request=account_request,
                            process=self._process)
    self.write(log_proto, flush, timestamp)

  def write_pnl_balance(self, pnl_balance, quote, subreq_id, *,
                        flush: bool = True, timestamp=None):
    assert isinstance(pnl_balance, CurrencyBalance)
    assert quote is None or isinstance(quote, CurrencyBalance)
    log_proto = StrategyLog(type=StrategyLog.PNL_BALANCE,
                            pnl_balance=pnl_balance,
                            process=self._process)
    self.write(log_proto, flush, timestamp)

  def write_pnl_balance_result(self, pnl_balance_result, *,
                               flush: bool = True, timestamp=None):
    assert isinstance(pnl_balance_result, StrategyLog)
    self.write(pnl_balance_result, flush, timestamp)

  def write_extra_info(self, extra_info: str, *,
                       flush: bool = True, timestamp=None):
    log_proto = StrategyLog(type=StrategyLog.EXTRA_INFO, extra_info=extra_info)
    self.write(log_proto, flush, timestamp)

  def write_og_log(self, og_log, *,
                   flush: bool = True, timestamp=None):
    assert isinstance(og_log, OrderGatewayLog)
    if og_log.type == OrderGatewayLog.BALANCE:
      prev = self._balance_cache
      dt = datetime.datetime.now()
      if dt.minute == 00:
        # For whole hour first minute, write anyway.
        pass
      else:
        if prev == og_log:
          return
      self._balance_cache = og_log

    if og_log.type == OrderGatewayLog.POSITION:
      prev = self._position_cache
      if prev == og_log:
        return
      self._position_cache = og_log

    log_proto = StrategyLog(type=StrategyLog.OG_LOG, og_log=og_log, process=self._process)
    self.write(log_proto, flush, timestamp)

  def write_request(self, *,
                    flush: bool = True, timestamp=None):
    log_proto = StrategyLog(type=StrategyLog.STRAT_REQUEST,
                            request=self.request,
                            process=self._process)
    self.write(log_proto, flush, timestamp)

  def write_executor_config(self, executor_config, *,
                            account_request=None, flush: bool = True, timestamp=None):
    assert account_request is None or \
        isinstance(account_request, AccountRequestProto)
    if isinstance(executor_config, AggressiveExecutorConfig):
      log_proto = StrategyLog(type=StrategyLog.AGG_EXEUCTOR_CONFIG,
                              account_request=account_request,
                              agg_executor_config=executor_config,
                              process=self._process)
    elif isinstance(executor_config, PassiveExecutorConfig):
      log_proto = StrategyLog(type=StrategyLog.PASS_EXECUTOR_CONFIG,
                              account_request=account_request,
                              pass_executor_config=executor_config,
                              process=self._process)
    else:
      raise ValueError('Unknown executor config: %s' % executor_config)
    self.write(log_proto, flush, timestamp)

  def write_best_book_entry(self, best_book_entry, *,
                            account_request=None, flush: bool = True, timestamp=None):
    assert isinstance(best_book_entry, BestBookEntry)
    assert account_request is None or \
        isinstance(account_request, AccountRequestProto)
    log_proto = StrategyLog(type=StrategyLog.BEST_BOOK,
                            best_book_entry=best_book_entry,
                            account_request=account_request,
                            process=self._process)
    self.write(log_proto, flush, timestamp)
